package de.szut.tree;

import java.util.Comparator;
import java.util.Iterator;
import java.util.Map.Entry;
import java.util.NoSuchElementException;

/**
 * Self balancing binary search tree
 * Offers the very basic tools for node manipulation
 * and many different interfaces for entry traversal.
 * @author Marc Huisinga
 *
 * @param <K> - Type of keys in entries in this tree
 * @param <V> - Type of values in entries in this tree
 */
public class AVLTree<K, V> implements Generator<Entry<K, V>> {  
	
	private static final String NULL_KEY = "Key must not be null.";
	private static final String KEY_NOT_FOUND = "Key not found.";
	private static final String NULL_START = "Start must not be null.";

	private Node<K, V> root = null;
	private int size = 0;

	private Comparator<? super K> comparator;

	/**
	 * Constructor
	 * @param comparator - comparator to order the elements in this tree with, must not be null
	 */
	public AVLTree(Comparator<? super K> comparator) {
		if (comparator == null) {
			throw new IllegalArgumentException("Comparator must not be null.");
		}
		this.comparator = comparator;
	}
	
	/**
	 * Copy constructor for this tree
	 * @param tree - tree to clone
	 */
	public AVLTree(AVLTree<K, V> tree) {
		if (tree == null) {
			throw new IllegalArgumentException("Tree must not be null.");
		}
		this.comparator = tree.comparator;
		tree.entries(Traversal.BREADTH_FIRST).forEach(entry -> put(entry.getKey(), entry.getValue()));
	}
	
	/**
	 * Iterates over the nodes in the direction of a key.
	 * @param start - node to start iterating at inclusively
	 * @param directionProvider - key to iterate in the direction of
	 * @return generator for lazy evaluation
	 */
	private Generator<Node<K, V>> nodesInKeyDirection(Node<K, V> start, K directionProvider) {
		return () -> {
			return new Iterator<Node<K, V>>() {

				Node<K, V> currentNode = start;
				
				@Override
				public boolean hasNext() {
					return currentNode != null;
				}

				@Override
				public Node<K, V> next() {
					if (!hasNext()) {
						throw new NoSuchElementException("End has been reached.");
					}
					K currentKey = currentNode.getEntry().getKey();
					int comparisonResult = comparator.compare(currentKey, directionProvider);
					Node<K, V> previousNode = currentNode;
					if (comparisonResult > 0) {
						currentNode = currentNode.getLeftChild();
					} else {
						currentNode = currentNode.getRightChild();
					}
					return previousNode;
				}
				
			};
		};
	}
	
	/**
	 * Replaces a node with another node.
	 * @param replaced - node to replace
	 * @param replacer - node to replace with
	 */
	private void replaceNode(Node<K, V> replaced, Node<K, V> replacer) {
		Node<K, V> parent = replaced.getParent();
		if (replaced != root) {
			parent.setMatchingChild(replaced, replacer);
		} else {
			root = replacer;
		}
		if (replacer != null) {
			replacer.setParent(parent);
		}
	}
	
	/**
	 * Rotates the specified node to the right.
	 * @param node - node to rotate
	 */
	private void rotateRight(Node<K, V> node) {
		Node<K, V> leftChild = node.getLeftChild();
		replaceNode(leftChild, leftChild.getRightChild());
		replaceNode(node, leftChild);
		leftChild.setRightChild(node);
		node.setParent(leftChild);
		node.updateHeight();
		leftChild.updateHeight();
	}
	
	/**
	 * Rotates the specified node to the left.
	 * @param node - node to rotate
	 */
	private void rotateLeft(Node<K, V> node) {
		Node<K, V> rightChild = node.getRightChild();
		replaceNode(rightChild, rightChild.getLeftChild());
		replaceNode(node, rightChild);
		rightChild.setLeftChild(node);
		node.setParent(rightChild);
		node.updateHeight();
		rightChild.updateHeight();
	}
	
	/**
	 * Balances the branch of a node that is at the edge of this tree.
	 * Since this tree is an AVL tree and all operations require balancing as well as height updating,
	 * this method does the height updating as well.
	 * Relies on the fact that each branch is being rebalanced after each imbalance 
	 * and that the balanceFactor() of each node in each branch is -2 <= n <= 2.
	 * Does not have perfect performance as it always checks the entire branch
	 * to be simple and usable for both .put() and .remove().
	 * @param edgeNode - node at a lower edge of this tree
	 */
	private void balanceBranch(Node<K, V> edgeNode) {
		IterationDirection.UPWARDS.iterate(edgeNode).forEach(node -> {
			node.updateHeight();
			int nodeBalance = node.getBalance();
			if (nodeBalance == 2) {
				Node<K, V> leftChild = node.getLeftChild();
				if (leftChild != null && leftChild.getBalance() == -1) {
					rotateLeft(leftChild);
				}
				rotateRight(node);
			} else if (nodeBalance == -2) {
				Node<K, V> rightChild = node.getRightChild();
				if (rightChild != null && rightChild.getBalance() == 1) {
					rotateRight(rightChild);
				}
				rotateLeft(node);
			}
		});
	}
	
	/**
	 * Puts an entry into this AVL tree.
	 * The value of the key is being updated should the key already be contained in the tree.
	 * @param key - key of the entry to insert, must not be null
	 * @param value - value of the entry to insert
	 * @return previous value
	 */
	public Replaced<V> put(K key, V value) {
		if (key == null) {
			throw new IllegalArgumentException(NULL_KEY);
		}
		if (root == null) {
			root = new Node<>(new TreeEntry<>(key, value));
			size++;
			return new Replaced<>();
		}
		Node<K, V> last = null;
		for (Node<K, V> node : nodesInKeyDirection(root, key)) {
			Entry<K, V> entry = node.getEntry();
			if (comparator.compare(entry.getKey(), key) == 0) {
				V old = entry.getValue();
				entry.setValue(value);
				return new Replaced<>(old);
			}
			last = node;
		}
		int comparisonResult = comparator.compare(last.getEntry().getKey(), key);
		Node<K, V> insertedNode = new Node<>(new TreeEntry<>(key, value), last);
		if (comparisonResult > 0) {
			last.setLeftChild(insertedNode);
		} else {
			last.setRightChild(insertedNode);
		}
		balanceBranch(last);
		size++;
		return new Replaced<>();
		
	}
	
	/**
	 * Finds the node of a specific key.
	 * @param key - key to find the node of
	 * @return node of the specific key or null if the key wasn't found
	 */
	private Node<K, V> find(K key) {
		return nodesInKeyDirection(root, key).stream()
				.filter(node -> comparator.compare(node.getEntry().getKey(), key) == 0)
				.findAny()
				.orElse(null);
	}
	
	/**
	 * Removes an entry from this tree.
	 * @param key - key of the entry to remove, must not be null and has to exist in this tree
	 * @return removed value
	 */
	public V remove(K key) {
		if (key == null) {
			throw new IllegalArgumentException(NULL_KEY);
		}
		Node<K, V> nodeToRemove = find(key);
		if (nodeToRemove == null) {
			throw new NoSuchElementException(KEY_NOT_FOUND);
		}
		V valueToRemove = nodeToRemove.getEntry().getValue();
		Node<K, V> leftChild = nodeToRemove.getLeftChild();
		Node<K, V> rightChild = nodeToRemove.getRightChild();
		if (leftChild != null && rightChild != null) {
			Node<K, V> nextLower = IterationDirection.RIGHT.iterate(leftChild).stream()
					.reduce((previous, current) -> current)
					.get();
			nodeToRemove.setEntry(nextLower.getEntry());
			nodeToRemove = nextLower;
			leftChild = nextLower.getLeftChild();
			rightChild = nextLower.getRightChild();
		}
		Node<K, V> parent = nodeToRemove.getParent();
 		if (leftChild != null) {
			replaceNode(nodeToRemove, leftChild);
		} else if (rightChild != null) {
			replaceNode(nodeToRemove, rightChild);
		} else {
			replaceNode(nodeToRemove, null);
		}
 		balanceBranch(parent);
 		size--;
		return valueToRemove;
	}
	
	/**
	 * Clears this tree.
	 */
	public void clear() {
		traverse(Traversal.POST_ORDER).stream().forEach(node -> {
			node.setLeftChild(null);
			node.setRightChild(null);
		});
		root = null;
		size = 0;
	}
	
	/**
	 * Gets the value mapped to a key.
	 * @param key - key to get the value of, must not be null and has to exist in this tree
	 * @return value
	 */
	public V get(K key) {
		if (key == null) {
			throw new IllegalArgumentException(NULL_KEY);
		}
		Node<K, V> node = find(key);
		if (node == null) {
			throw new NoSuchElementException(KEY_NOT_FOUND);
		}
		return node.getEntry().getValue();
	}
	
	/**
	 * Checks if the specified key is contained in this tree.
	 * @param key - key to check, must not be null
	 * @return true if the tree is contained in this tree or false if not
	 */
	public boolean contains(K key) {
		if (key == null) {
			throw new IllegalArgumentException(NULL_KEY);
		}
		return find(key) != null ? true : false;
	}
	
	@SuppressWarnings({ "rawtypes", "unchecked" })
	@Override
	public boolean equals(Object obj) {
		if (this == obj) {
			return true;
		}
		if (obj == null) {
			return false;
		}
		if (getClass() != obj.getClass()) {
			return false;
		}
		AVLTree other = (AVLTree) obj;
		if (comparator == null) {
			if (other.comparator != null) {
				return false;
			}
		} else if (!comparator.equals(other.comparator)) {
			return false;
		}
		if (size != other.size) {
			return false;
		}
		Iterator<Node> otherNodes = Traversal.BREADTH_FIRST.traverse(other.root).iterator();
		boolean allEqual = Traversal.BREADTH_FIRST.traverse(root).stream()
				.allMatch(node -> {
					if (!otherNodes.hasNext()) {
						return false;
					}
					Node next = otherNodes.next();
					if (node == null) {
						if (next != null) {
							return false;
						}
					} else if (!node.equals(next)) {
						return false;
					}
					return true;
				});
		if (otherNodes.hasNext()) {
			return false;
		}
		if (!allEqual) {
			return false;
		}
		return true;
	}
	
	/**
	 * Converts a node generator to an entry generator.
	 * @param nodes - node generator to convert
	 * @return converted entry generator
	 */
	private EntryGenerator<K, V> nodesToEntries(Generator<Node<K, V>> nodes) {
		Generator<Entry<K, V>> entries = () -> nodes.stream()
				.map(Node::getEntry)
				.map(entry -> (Entry<K, V>) new TreeEntry<>(entry))
				.iterator();
		return new EntryGenerator<K, V>(entries);
	}
	
	/**
	 * Iterates the tree in a specific direction until an end of the tree is reached.
	 * The end of the tree is inclusive, meaning the last node before the end will be provided
	 * by this iteration.
	 * @param start - key to inclusively start iterating at, must not be null
	 * @param direction - direction to iterate in, must not be null
	 * @return generator for lazy evaluation, is empty if the specified start wasn't found
	 */
	public EntryGenerator<K, V> inDirection(K start, IterationDirection direction) {
		if (start == null) {
			throw new IllegalArgumentException(NULL_START);
		}
		if (direction == null) {
			throw new IllegalArgumentException("Direction must not be null.");
		}
		Node<K, V> startNode = find(start);
		if (startNode == null) {
			return EntryGenerator.empty();
		}
		return nodesToEntries(direction.iterate(startNode));
	}
	
	/**
	 * Iterates the tree in the direction of a key until an end of the tree is reached.
	 * The step in which the direction provider is passed over will result in an iteration towards the
	 * lower right entry of the direction provider, the right child.
	 * The end of the tree is inclusive, meaning the last node before the end will be provided
	 * by this iteration.
	 * @param start - key to inclusively start iterating at, must not be null
	 * @param directionProvider - key to specify the iteration direction, must not be null
	 * @return generator for lazy evaluation, is empty if the specified start wasn't found
	 */
	public EntryGenerator<K, V> inDirection(K start, K directionProvider) {
		if (start == null) {
			throw new IllegalArgumentException(NULL_START);
		}
		if (directionProvider == null) {
			throw new IllegalArgumentException("Direction provider must not be null.");
		}
		Node<K, V> startNode = find(start);
		if (startNode == null) {
			return EntryGenerator.empty();
		}
		return nodesToEntries(nodesInKeyDirection(startNode, directionProvider));
	}

	/**
	 * Traverses this tree with the specified traversal.
	 * @param traversal - traversal to traverse with
	 * @return generator for lazy evaluation
	 */
	private Generator<Node<K, V>> traverse(Traversal traversal) {
		return traversal.traverse(root);
	}

	/**
	 * Traverses the entries in this tree using the specified traversal order 
	 * until no elements are left to traverse.
	 * @param traversal - traversal order to traverse the tree with, must not be null
	 * @return generator for lazy evaluation
	 */
	public EntryGenerator<K, V> entries(Traversal traversal) {
		if (traversal == null) {
			throw new IllegalArgumentException("Traversal must not be null.");
		}
		return nodesToEntries(traverse(traversal));
	}
	
	@Override
	public Iterator<Entry<K, V>> iterator() {
		return entries(Traversal.ASCENDING).iterator();
	}
	
	/**
	 * Gets the comparator used to order the elements in this tree with.
	 * @return comparator
	 */
	public Comparator<? super K> comparator() {
		return comparator;
	}

	/**
	 * Gets the the amount of entries in this tree.
	 * @return size
	 */
	public int size() {
		return size;
	}

	/**
	 * Gets the the amount of vertical layers in this tree.
	 * @return height
	 */
	public int height() {
		return root.getHeight();
	}
	
	/**
	 * Gets the root entry of this tree.
	 * @return root entry
	 */
	public Entry<K, V> root() {
		return root == null ? null : new TreeEntry<>(root.getEntry());
	}
	
}
